package jam4j.util;

import jam4j.util.Throwables.ExceptionCapturer;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.io.OutputStream;
import java.util.*;
import java.util.concurrent.CountDownLatch;

/**
 * Spawns a number of threads, each of which handles continuously piping data
 * from one stream into another.
 * <p>
 * This class is not thread-safe, despite its asynchronous nature; only one
 * thread should manage an AsyncMultiPipe object.
 * 
 * @author Luke Maurer
 */
public final class AsyncMultiPipe {
    private final Pipe[] pipes;
    private final int count;
    
    private final ThreadGroup group;
    private final CountDownLatch start, finish;
    
    private boolean started;
    
    private final List<IOException> exceptions =
        Collections.synchronizedList(new ArrayList<IOException>());
    private final ExceptionCapturer<? super IOException> capturer =
        Throwables.captureInto(exceptions);
    
    /**
     * Begin building a multi-pipe.
     * 
     * @return A {@link Builder} object.
     * 
     * @see Builder
     */
    public static Builder build() {
        return new Builder();
    }
    
    /**
     * Builder for {@link AsyncMultiPipe} objects. All methods are optional,
     * though the multi-pipe won't do much without at least one addPipe call.
     * 
     * @see AsyncMultiPipe#build()
     * 
     * @author Luke Maurer
     */
    public static final class Builder {
        private Builder() { }
        
        private String name;
        private ThreadGroup group;
        private final List<Pipe> pipes = new ArrayList<Pipe>();
        
        /**
         * Set the name of the multi-pipe. This will serve as the name of its
         * thread group.
         * 
         * @param name The name to give the multi-pipe and its thread group.
         * @return This object.
         */
        public Builder name(String name) {
            this.name = name;
            return this;
        }
        
        /**
         * Specify the thread group under which to create the threads which
         * will handle the pipes. (Actually it will be the parent of the
         * thread <em>group</em> which will contain the threads.)
         * <p>
         * By default the current thread's group is used.
         * 
         * @param group The group to make the parent of the multi-pipe's thread
         *                group.
         * @return This object.
         */
        public Builder parentGroup(ThreadGroup group) {
            this.group = group;
            return this;
        }
        
        /**
         * Add a pipe to the multi-pipe. This will channel the input stream
         * into the output stream asynchronously, once the multi-pipe is
         * started.
         * 
         * @param in The input stream from which to read data.
         * @param out The output stream into which to write data.
         * @return This object.
         */
        public Builder addPipe(InputStream in, OutputStream out) {
            pipes.add(new Pipe(pipes.size(), in, out));
            return this;
        }
        
        /**
         * Add a pipe to the multi-pipe. This will channel the input stream
         * into the output stream asynchronously, once the multi-pipe is
         * started.
         * 
         * @param in The input stream from which to read data.
         * @param out The output stream into which to write data.
         * @param name The name to give this pipe. This will serve as the name
         *                of the thread which handles the pipe.
         * @return This object.
         */
        public Builder addPipe(InputStream in, OutputStream out, String name) {
            pipes.add(new Pipe(name, in, out));
            return this;
        }
        
        /**
         * Add pipes into and out of a process.
         * 
         * @param process The process to attach to.
         * @param in The input stream to feed into the process's standard
         *                input.
         * @param out The output stream into which to read from the process's
         *                standard output.
         * @param err The output stream into which to read from the process's
         *                standard error.
         * @param prefix The prefix to give to the names of the pipes.
         * @return This object.
         */
        public Builder addPipes(Process process, 
                InputStream in, OutputStream out,
                OutputStream err, String prefix) {
            if (in != null)
                addPipe(in, process.getOutputStream(), prefix + "in");
            addPipe(process.getInputStream(), out, prefix + "out");
            addPipe(process.getErrorStream(), err, prefix + "err");
            
            return this;
        }
        
        /**
         * Add pipes out of a process.
         * 
         * @param process The process to attach to.
         * @param out The output stream into which to read from the process's
         *                standard output.
         * @param err The output stream into which to read from the process's
         *                standard error.
         * @param prefix The prefix to give to the names of the pipes.
         * @return This object.
         */
        public Builder addPipes(Process process,
                OutputStream out, OutputStream err, String prefix) {
            return addPipes(process, null, out, err, prefix);
        }
        
        /**
         * Add pipes into and out of a process.
         * 
         * @param process The process to attach to.
         * @param in The input stream to feed into the process's standard
         *                input.
         * @param out The output stream into which to read from the process's
         *                standard output.
         * @param err The output stream into which to read from the process's
         *                standard error.
         * @return This object.
         */
        public Builder addPipes(Process process,
                InputStream in, OutputStream out, OutputStream err) {
            return addPipes(process, in, out, err, "");
        }
        
        /**
         * Add pipes out of a process.
         * 
         * @param process The process to attach to.
         * @param out The output stream into which to read from the process's
         *                standard output.
         * @param err The output stream into which to read from the process's
         *                standard error.
         * @return This object.
         */
        public Builder addPipes(Process process,
                OutputStream out, OutputStream err) {
            return addPipes(process, null, out, err, "");
        }
        
        /**
         * Get the built multi-pipe.
         * 
         * @return The result.
         */
        public AsyncMultiPipe done() {
            return new AsyncMultiPipe(group, name, pipes);
        }
    }
    
    /**
     * Create a multi-pipe between a process's input, output, and error streams
     * and the given ones.
     * 
     * @param process The process to attach to.
     * @param in The stream to pipe into the process's standard input.
     * @param out The stream into which to pipe the process's standard output. 
     * @param err The stream into which to pipe the process's standard error.
     * @return A multi-pipe which will pipe the streams as given.
     * 
     * @see Builder#addPipes(Process, InputStream, OutputStream, OutputStream)
     */
    public static AsyncMultiPipe fromProcess(Process process,
            InputStream in, OutputStream out, OutputStream err) {
        return build().addPipes(process, in, out, err).done();
    }
    
    /**
     * Create a multi-pipe from a process's output and error streams to the
     * given ones.
     * 
     * @param process The process to attach to.
     * @param out The stream into which to pipe the process's standard output.
     * @param err The stream into which to pipe the process's standard error.
     * @return A multi-pipe which will pipe the streams as given.
     * 
     * @see Builder#addPipes(Process, OutputStream, OutputStream)
     */
    public static AsyncMultiPipe fromProcess(Process process,
            OutputStream out, OutputStream err) {
        return build().addPipes(process, out, err).done();
    }
    
    private AsyncMultiPipe(ThreadGroup group, String name, List<Pipe> pipes) {
        if (group == null)
            group = Thread.currentThread().getThreadGroup();
        if (name == null)
            name = "pipes";
        
        this.pipes = pipes.toArray(new Pipe[this.count = pipes.size()]);
        this.start = new CountDownLatch(1);
        this.finish = new CountDownLatch(this.count);
        this.group = new ThreadGroup(group, name);
    }
    
    /**
     * Begin piping data, and do so until all input streams are exhausted.
     * 
     * @throws IOException If any stream throws it.
     * @throws InterruptedException If an I/O operation is interrupted.
     * 
     * @see #fork()
     * @see #join()
     * @see #throwExceptions()
     */
    public void pipe()
            throws IOException, InterruptedException {
        fork();
        join();
        throwExceptions();
    }
    
    /**
     * Begin piping data in the background.
     */
    public void fork() {
        for (int ix = 0; ix < count; ix++)
            new PipeThread(pipes[ix]).start();
        
        started = true;
        start.countDown();
    }
    
    /**
     * Wait until all input streams are exhausted and all output streams have
     * been written to.
     * 
     * @throws IllegalStateException If {@link #fork()} has not been called.
     * @throws InterruptedException If the thread, or an I/O operation, is
     * interrupted.
     */
    public void join() throws InterruptedException {
        if (!started)
            throw new IllegalStateException("Not started");
        
        try {
            finish.await();
        } catch (InterruptedException e) {
            group.interrupt();
            
            throw e;
        }
        
        started = false;
    }
    
    /**
     * Throw the exceptions that have been thrown by the pipe threads and
     * captured asynchronously. Multiple exceptions will be bundled using a
     * {@link MultipleIOExceptions} object.
     * 
     * @throws IOException If one has been thrown by a pipe thread.
     */
    public void throwExceptions() throws IOException {
        synchronized (exceptions) {
            try {
                MultipleIOExceptions.throwAll(exceptions);
            } finally {
                exceptions.clear();
            }
        }
    }
    
    private static final class Pipe {
        final InputStream in;
        final OutputStream out;
        final String name;
        
        Pipe(String name, InputStream in, OutputStream out) {
            this.name = name;
            this.in = in;
            this.out = out;
        }
        
        Pipe(int index, InputStream in, OutputStream out) {
            this("pipe-" + index, in, out);
        }
    }
    
    private final class PipeThread extends Thread {
        private final InputStream in;
        private final OutputStream out;
        
        PipeThread(Pipe pipe) {
            super(group, pipe.name);
            this.in = pipe.in;
            this.out = pipe.out;
            
            this.setDaemon(true);
        }
        
        public void run() {
            try {
                try {
                    start.await();
                } catch (InterruptedException e) {
                    return;
                }
                
                try {
                    while (!interrupted()) {
                        final int b;
                        if ((b = in.read()) == -1)
                            return;
                        out.write(b);
                    }
                } catch (IOException e) {
                    capturer.capture(e);
                }
            } finally {
                finish.countDown();
            }
        }
    }
    
    /**
     * Capture all data from each of the given input streams into a byte array,
     * returning the byte arrays when complete.
     * 
     * @param ins The streams to capture.
     * @return The data captured from each stream, as a byte array.
     * @throws IOException If any stream throws it.
     * @throws InterruptedIOException If a thread is interrupted.
     */
    public static byte[][] capture(InputStream ... ins) throws IOException {
        final int count = ins.length;
        
        final Builder builder = build();
        final ByteArrayOutputStream[] outs = new ByteArrayOutputStream[count];
        for (int ix = 0; ix < count; ix++)
            builder.addPipe(ins[ix], outs[ix] = new ByteArrayOutputStream());
        final AsyncMultiPipe multiPipe = builder.done();
        
        try {
            multiPipe.pipe();
        } catch (InterruptedException e) {
            throw new InterruptedIOException();
        }
        
        final byte[][] ans = new byte[count][];
        for (int ix = 0; ix < count; ix++)
            ans[ix] = outs[ix].toByteArray();
        
        return ans;
    }
    
    /**
     * Capture data from the standard output and error streams of the given
     * process, returning two byte arrays when the process finishes.
     * 
     * @param process The process to capture output from.
     * @return A two-element array with the output data and the error data.
     * @throws IOException If any stream throws it.
     * @throws InterruptedIOException If a thread is interrupted.
     */
    public static byte[][] capture(Process process) throws IOException {
        final InputStream
            in = process.getInputStream(),
            errIn = process.getErrorStream();
        
        try {
            return capture(in, errIn);
        } finally {
            try {
                in.close();
            } catch (IOException e) { }
            
            try {
                errIn.close();
            } catch (IOException e) { }
        }
    }
}
